import argparse
import netaddr
import os
import re
import requests
import sys
import time
import pymysql.cursors

from requests import HTTPError
from f5.bigip import ManagementRoot
from requests.packages.urllib3.exceptions import InsecureRequestWarning

from keystoneauth1 import identity
from keystoneauth1 import session
from neutronclient.v2_0 import client

requests.packages.urllib3.disable_warnings(InsecureRequestWarning)


uuid_regex = "[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"    # noqa

partition_regex = "^[a-zA-Z0-9]+_[a-z0-9]{32}$"

vip_regex = "^[a-zA-Z0-9]+_" + uuid_regex + "$"

vs_regex = vip_regex

snat_pool_regex = partition_regex

map = {
    "bigip": {},
    "neutron": {
        "network": {},
        "subnet": {},
        "port": {}
    }
}

flavor_map = {
    4: {1: 2, 2: 2, 3: 2, 4: 4, 5: 8, 6: 16},
    6: {1: 1, 2: 1, 3: 1, 4: 3, 5: 7, 6: 15}
}

device_owner = "network:f5lbaasv2"


class Context(object):

    def __init__(self):
        self.db = None
        self.bigips = []
        self.neutron = None


def loop_in_bigips(func):
    def loop(bigips, *args):
        if type(bigips) is list:
            for bigip in bigips:
                func(bigip, *args)
        else:
            func(bigips, *args)
    return loop


def retry_http_errors(*codes, **params):
    max_retry = params.get("retry", 1)
    interval = params.get("interval", 0)

    def retry(func):

        def wrapper(*args, **kwargs):
            i = 0
            while i < max_retry + 1:
                try:
                    i = i + 1
                    return func(*args, **kwargs)
                except HTTPError as ex:
                    need_to_retry = False
                    for code in codes:
                        if code == ex.response.status_code:
                            need_to_retry = True
                            break
                    if need_to_retry and i < max_retry + 1:
                        print("Retry on HTTP error %s", code)
                        if interval > 0:
                            time.sleep(interval)
                    elif i >= max_retry + 1:
                        print("Exceed HTTP errors retry limit")
                        raise ex
                    else:
                        raise ex

        return wrapper

    return retry


@retry_http_errors(401)
def retry_icontrol(icontrol, **kwargs):
    return icontrol(**kwargs)


@loop_in_bigips
def get_device_name(bigip):
    devices = bigip.tm.cm.devices.get_collection()
    for device in devices:
        if device.selfDevice == "true":
            bigip.device_name = device.name

    if not hasattr(bigip, "device_name"):
        print >> sys.stderr, "Cannot get device name of BIG-IP " + \
            bigip.hostname
        sys.exit(1)


def get_neutron_client():
    username = os.environ['OS_USERNAME']
    password = os.environ['OS_PASSWORD']
    project_name = os.environ['OS_PROJECT_NAME']
    project_domain_id = os.environ['OS_PROJECT_DOMAIN_ID']
    user_domain_id = os.environ['OS_USER_DOMAIN_ID']
    auth_url = os.environ['OS_AUTH_URL']
    auth = identity.Password(auth_url=auth_url,
                             username=username,
                             password=password,
                             project_name=project_name,
                             project_domain_id=project_domain_id,
                             user_domain_id=user_domain_id)
    sess = session.Session(auth=auth)
    return client.Client(session=sess)


def create_snat_port(neutron):
    print "Start creating dedicated SNAT port"
    subnet = map["neutron"]["subnet"]
    port = map["neutron"]["port"]

    for addr in map["bigip"].keys():
        partition = map["bigip"][addr]["partition"]
        for par_name in partition.keys():
            par = partition[par_name]
            for vip_name in par["virtual_address"].keys():
                vip = par["virtual_address"][vip_name]
                lb_id = vip_name[vip_name.rfind("_")+1:]
                vip_subnet_id = vip["subnet_id"]
                network_id = subnet[vip_subnet_id]["network_id"]

                v4_subnet_id = None
                v6_subnet_id = None
                v4_ip_count = 0
                v6_ip_count = 0
                for id in subnet.keys():
                    if subnet[id]["network_id"] == network_id:
                        if subnet[id]["ip_version"] == 4:
                            v4_subnet_id = id
                            v4_ip_count = flavor_map[4][vip["flavor"]]
                        if subnet[id]["ip_version"] == 6:
                            v6_subnet_id = id
                            v6_ip_count = flavor_map[6][vip["flavor"]]

                port["snat_ipv4_" + lb_id] = {
                    "tenant_id": par_name[-32:],
                    "network_id": network_id,
                    "subnet_id": v4_subnet_id,
                    "ip_count": v4_ip_count
                }

                port["snat_ipv6_" + lb_id] = {
                    "tenant_id": par_name[-32:],
                    "network_id": network_id,
                    "subnet_id": v6_subnet_id,
                    "ip_count": v6_ip_count,
                }

    resp = neutron.list_ports(device_owner=device_owner)
    legacy_ports = resp['ports']

    for port_name in port.keys():
        # Needn't to create this snat port
        if port[port_name]["ip_count"] <= 0:
            continue

        # Port already exists, skip to create it
        resp = neutron.list_ports(name=port_name)
        if len(resp["ports"]) > 0:
            port[port_name]["fixed_ips"] = resp["ports"][0]["fixed_ips"]
            continue

        legacy = legacy_ports[0]
        subnet_id = port[port_name]["subnet_id"]
        network_id = port[port_name]["network_id"]
        for p in legacy_ports:
            if subnet_id in p["name"]:
                legacy = p

        body = {
            "port": {
                "name": port_name,
                "tenant_id": port[port_name]["tenant_id"],
                "network_id": network_id,
                "device_id": port_name[-36:],
                "device_owner": device_owner,
                "port_security_enabled": False,
                "binding:vnic_type": legacy["binding:vnic_type"],
                "binding:host_id": legacy["binding:host_id"],
                "fixed_ips": []
            }
        }
        for i in range(port[port_name]["ip_count"]):
            body["port"]["fixed_ips"].append({"subnet_id": subnet_id})

        print "Creating SNAT port " + port_name
        resp = neutron.create_port(body)
        port[port_name]["fixed_ips"] = resp["port"]["fixed_ips"]

    print "Complete creating dedicated SNAT port"


def delete_snat_port(neutron):
    print "Start deleting dedicated SNAT port"

    for addr in map["bigip"].keys():
        partition = map["bigip"][addr]["partition"]
        for par_name in partition.keys():
            par = partition[par_name]
            for vip_name in par["virtual_address"].keys():
                lb_id = vip_name[vip_name.rfind("_")+1:]
                for port_name in ["snat_ipv4_" + lb_id, "snat_ipv6_" + lb_id]:
                    print "Deleting SNAT port " + port_name
                    delete_ports_by_name(neutron, port_name)

    print "Complete deleting dedicated SNAT port"


def delete_legacy_snat_port(neutron):
    print "Start deleting legacy SNAT port"

    resp = neutron.list_ports(device_owner=device_owner)
    ports = resp['ports']
    for port in ports:
        if port["name"].startswith("snat-") and len(port["name"]) > 41:
            print "Deleting legacy SNAT port " + port["name"]
            neutron.delete_port(port["id"])

    print "Complete deleting legacy SNAT port"


def delete_ports_by_name(neutron, name):
    resp = neutron.list_ports(name=name)
    ports = resp["ports"]
    for port in ports:
        neutron.delete_port(port["id"])


def connect_neutron_db(db_conn):
    db_type = db_conn[0:db_conn.find("+")]
    if db_type != "mysql":
        print >> sys.stderr, "Only support MySQL DB"
        sys.exit(1)

    u_start = db_conn.find("//") + 2
    u_end = db_conn.find(":", u_start)
    p_start = u_end + 1
    p_end = db_conn.find("@", p_start)
    h_start = p_end + 1
    h_end = db_conn.find("/", h_start)
    d_start = h_end + 1

    return pymysql.connect(host=db_conn[h_start:h_end],
                           user=db_conn[u_start:u_end],
                           password=db_conn[p_start:p_end],
                           database=db_conn[d_start:],
                           cursorclass=pymysql.cursors.DictCursor)


def load_loadbalancer(connection, id):
    with connection:
        with connection.cursor() as cursor:
            sql = ("SELECT `id`, `flavor`, `vip_subnet_id` "
                   "FROM `lbaas_loadbalancers` "
                   "WHERE `id`=%s")
            cursor.execute(sql, (id))
            return cursor.fetchone()


def load_subnet_and_network(conn, subnet_id):
    subnet = load_subnet(conn, subnet_id)
    if not subnet:
        print >> sys.stderr, "Fail to load subnet " + \
            subnet_id + " from db"
        return
    net_id = subnet["network_id"]
    map["neutron"]["subnet"][subnet_id] = {
        "cidr": subnet["cidr"],
        "ip_version": subnet["ip_version"],
        "network_id": net_id
    }
    if net_id not in map["neutron"]["network"]:
        net = load_network(conn, net_id)
        if not net:
            print >> sys.stderr, "Fail to load network" + \
                net_id + " from db"
            return
        map["neutron"]["network"][net_id] = {
            "segmentation_id": net["segmentation_id"]
        }
        subnets = load_subnet_by_network(conn, net_id)
        for s in subnets:
            if s["id"] not in map["neutron"]["subnet"]:
                map["neutron"]["subnet"][s["id"]] = {
                    "cidr": s["cidr"],
                    "ip_version": s["ip_version"],
                    "network_id": net_id
                }


def load_network(connection, id):
    with connection:
        with connection.cursor() as cursor:
            sql = """
            SELECT net.id, seg.segmentation_id
            FROM networks net
            INNER JOIN networksegments seg
            ON net.id = seg.network_id
            WHERE net.id=%s AND seg.network_type='vlan'
            """
            cursor.execute(sql, (id))
            return cursor.fetchone()


def load_subnet(connection, id):
    with connection:
        with connection.cursor() as cursor:
            sql = """
            SELECT `id`, `cidr`, `ip_version`, `network_id`
            FROM `subnets`
            WHERE `id`=%s
            """
            cursor.execute(sql, (id))
            return cursor.fetchone()


def load_subnet_by_network(connection, id):
    with connection:
        with connection.cursor() as cursor:
            sql = """
            SELECT `id`, `cidr`, `ip_version`, `network_id`
            FROM `subnets`
            WHERE `network_id`=%s
            """
            cursor.execute(sql, (id))
            return cursor.fetchall()


@loop_in_bigips
def load_partition(bigip):
    addr = bigip.hostname
    map["bigip"][addr] = {}
    map["bigip"][addr]["partition"] = {}
    partition = map["bigip"][addr]["partition"]
    pattern = re.compile(partition_regex)

    pars = retry_icontrol(bigip.tm.sys.folders.get_collection)

    for par in pars:
        if re.match(pattern, par.name):
            partition[par.name] = {
                "old_snat_pool": {},
                "old_snat_translation": {},
                "virtual_address": {},
                "virtual_server": {}
            }


@loop_in_bigips
def load_vip(bigip):
    addr = bigip.hostname
    partition = map["bigip"][addr]["partition"]
    vips = retry_icontrol(bigip.tm.ltm.virtual_address_s.get_collection)
    pattern = re.compile(vip_regex)
    for vip in vips:
        if vip.partition in partition.keys() and \
           re.match(pattern, vip.name):
            par = partition[vip.partition]
            par["virtual_address"][vip.name] = {
                "address": vip.address
            }


@loop_in_bigips
def load_vs(bigip):
    addr = bigip.hostname
    partition = map["bigip"][addr]["partition"]
    vs_s = retry_icontrol(bigip.tm.ltm.virtuals.get_collection)
    pattern = re.compile(vs_regex)
    for vs in vs_s:
        if vs.partition in partition.keys() and \
           re.match(pattern, vs.name):
            par = partition[vs.partition]
            sat = vs.sourceAddressTranslation
            if "poolReference" in sat:
                del sat["poolReference"]
            par["virtual_server"][vs.name] = {
                "destination": vs.destination,
                "sourceAddressTranslation": sat
            }


def load_lbaas_db(conn):
    for addr in map["bigip"].keys():
        partition = map["bigip"][addr]["partition"]
        for par_name in partition.keys():
            par = partition[par_name]

            # Load loadbalancer flavor and subnet_id
            for vip_name in par["virtual_address"].keys():
                vip = par["virtual_address"][vip_name]
                lb_id = vip_name[vip_name.rfind("_")+1:]
                lb = load_loadbalancer(conn, lb_id)
                if not lb:
                    print >> sys.stderr, "Fail to load loadbalancer " + \
                        lb_id + " from db, whose vip in bigip is " + \
                        "/" + par_name + "/" + vip_name
                    continue
                subnet_id = lb["vip_subnet_id"]
                vip["subnet_id"] = subnet_id
                vip["flavor"] = lb["flavor"]
                # Load subnet and network
                if subnet_id not in map["neutron"]["subnet"]:
                    load_subnet_and_network(conn, subnet_id)


@loop_in_bigips
def create_dedicated_snat_pool(bigip):
    addr = bigip.hostname
    partition = map["bigip"][addr]["partition"]
    network = map["neutron"]["network"]
    port = map["neutron"]["port"]
    snatpool = bigip.tm.ltm.snatpools.snatpool
    print "Start creating dedicated SNAT pool on BIG-IP " + addr

    for par_name in partition.keys():
        par = partition[par_name]

        pool = {}
        for vip in par["virtual_address"].keys():
            lb_id = vip[-36:]
            pool[vip] = {
                "members": []
            }
            for port_name in ["snat_ipv4_" + lb_id, "snat_ipv6_" + lb_id]:
                if "fixed_ips" in port[port_name]:
                    for ip in port[port_name]["fixed_ips"]:
                        seg = network[
                            port[port_name]["network_id"]
                        ]["segmentation_id"]
                        pool[vip]["members"].append(
                            ip["ip_address"] + "%" + str(seg)
                        )

        for pool_name in pool.keys():
            payload = {
                "partition": par_name,
                "name": pool_name,
                "members": pool[pool_name]["members"]
            }
            # Skip to create dedicated snat pool, if it already exists
            if not retry_icontrol(snatpool.exists, partition=par_name,
                                  name=pool_name):
                print "Creating SNAT pool " + pool_name
                retry_icontrol(snatpool.create, **payload)

    print "Complete creating dedicated SNAT pool on BIG-IP " + addr


@loop_in_bigips
def delete_dedicated_snat_pool(bigip):
    addr = bigip.hostname
    partition = map["bigip"][addr]["partition"]
    snatpool = bigip.tm.ltm.snatpools.snatpool
    print "Start deleting dedicated SNAT pool on BIG-IP " + addr

    for par_name in partition.keys():
        par = partition[par_name]

        pool = {}
        for vip_name in par["virtual_address"].keys():
            if retry_icontrol(snatpool.exists, partition=par_name,
                              name=vip_name):
                pool = retry_icontrol(snatpool.load, partition=par_name,
                                      name=vip_name)
                print "Deleting dedicated SNAT pool " + vip_name
                retry_icontrol(pool.delete)

    print "Complete deleting dedicated SNAT pool on BIG-IP " + addr


@loop_in_bigips
def switch_to_dedicated_snat_pool(bigip):
    addr = bigip.hostname
    partition = map["bigip"][addr]["partition"]
    virtual = bigip.tm.ltm.virtuals.virtual

    for par_name in partition.keys():
        par = partition[par_name]
        for vs_name in par["virtual_server"].keys():
            vs = par["virtual_server"][vs_name]
            if vs["sourceAddressTranslation"]["type"] == "snat":
                dest = vs["destination"]
                vip = dest[:dest.rfind(":")]
                payload = {
                    "snatpool": vip
                }
                v = retry_icontrol(virtual.load, partition=par_name,
                                   name=vs_name)
                if v.sourceAddressTranslation["pool"] != payload["snatpool"]:
                    print "Switching SNAT pool of VS " + vs_name
                    retry_icontrol(v.modify, **payload)


@loop_in_bigips
def switch_to_shared_snat_pool(bigip):
    addr = bigip.hostname
    partition = map["bigip"][addr]["partition"]
    virtual = bigip.tm.ltm.virtuals.virtual

    for par_name in partition.keys():
        par = partition[par_name]
        for vs_name in par["virtual_server"].keys():
            vs = par["virtual_server"][vs_name]
            if vs["sourceAddressTranslation"]["type"] == "snat":
                payload = {
                    "snatpool": "/" + par_name + "/" + par_name
                }
                v = retry_icontrol(virtual.load, partition=par_name,
                                   name=vs_name)
                if v.sourceAddressTranslation["pool"] != payload["snatpool"]:
                    print "Switching SNAT pool of VS " + vs_name
                    retry_icontrol(v.modify, **payload)


@loop_in_bigips
def delete_shared_snat_pool(bigip):
    addr = bigip.hostname
    partition = map["bigip"][addr]["partition"]
    snatpool = bigip.tm.ltm.snatpools.snatpool
    print "Start deleting shared SNAT pool on BIG-IP " + addr

    for par_name in partition.keys():
        if retry_icontrol(snatpool.exists, partition=par_name, name=par_name):
            pool = retry_icontrol(snatpool.load, partition=par_name,
                                  name=par_name)
            print "Deleting shared SNAT pool " + par_name
            retry_icontrol(pool.delete)

    print "Complete deleting shared SNAT pool on BIG-IP " + addr


@loop_in_bigips
def create_snat_selfip(bigip, neutron):
    addr = bigip.hostname
    partition = map["bigip"][addr]["partition"]
    network = map["neutron"]["network"]
    subnet = map["neutron"]["subnet"]
    selfip = bigip.tm.net.selfips.selfip
    selfips = bigip.tm.net.selfips
    print "Start creating SelfIP"

    all_selfips = retry_icontrol(selfips.get_collection)

    resp = neutron.list_ports(device_owner=device_owner)
    legacy_ports = resp['ports']

    snat_subnet = subnet.keys()

    for par_name in partition.keys():
        par = partition[par_name]
        for vip_name in par["virtual_address"].keys():
            vip = par["virtual_address"][vip_name]
            vip_subnet_id = vip["subnet_id"]
            try:
                snat_subnet.remove(vip_subnet_id)
            except ValueError:
                pass

    for subnet_id in snat_subnet:
        network_id = subnet[subnet_id]["network_id"]
        legacy = None
        selfip_par = None
        for legacy_port in legacy_ports:
            if legacy_port["network_id"] == network_id and \
               legacy_port["name"].startswith("local-" + bigip.device_name):
                legacy = legacy_port
                for par_name in partition.keys():
                    if par_name.endswith(legacy["tenant_id"]):
                        selfip_par = par_name
                break

        if not legacy:
            print >> sys.stderr, "Cannot find SelfIP port in network " + \
                network_id
            sys.exit(1)

        selfip_port_name = "local-" + bigip.device_name + "-" + subnet_id

        skip_it = False
        for s in all_selfips:
            if s.name == selfip_port_name:
                skip_it = True
                break

        if skip_it:
            continue

        body = {
            "port": {
                "name": selfip_port_name,
                "tenant_id": legacy["tenant_id"],
                "network_id": network_id,
                "device_id": legacy["device_id"],
                "device_owner": device_owner,
                "port_security_enabled": False,
                "binding:vnic_type": legacy["binding:vnic_type"],
                "binding:host_id": legacy["binding:host_id"],
                "fixed_ips": [{"subnet_id": subnet_id}]
            }
        }

        print "Creating SelfIP port " + body["port"]["name"]
        resp = neutron.create_port(body)
        seg_id = str(network[network_id]["segmentation_id"])
        mask = str(netaddr.IPNetwork(subnet[subnet_id]["cidr"]).prefixlen)
        selfip_addr = resp["port"]["fixed_ips"][0]["ip_address"] + \
            "%" + seg_id + "/" + mask
        vlan = "vlan-" + seg_id

        selfip_body = {
            "name": selfip_port_name,
            "address": selfip_addr,
            "vlan": vlan,
            "floating": "disabled",
            "partition": selfip_par
        }
        print "Creating SelfIP " + selfip_addr + " in BIG-IP " + \
              bigip.hostname + " under partition " + selfip_par
        retry_icontrol(selfip.create, **selfip_body)

    print "Complete creating SelfIP"


@loop_in_bigips
def delete_snat_selfip(bigip):
    pass


@loop_in_bigips
def tmsh_save(bigip):
    retry_icontrol(bigip.tm.util.bash.exec_cmd,
                   command="run",
                   utilCmdArgs="-c 'tmsh save sys config'")


def scan_network():
    network = map["neutron"]["network"]
    subnet = map["neutron"]["subnet"]

    ret = 0
    for network_id in network.keys():
        v4_count = 0
        v6_count = 0
        for subnet_id in subnet.keys():
            if subnet[subnet_id]["network_id"] == network_id:
                if subnet[subnet_id]["ip_version"] == 4:
                    v4_count = v4_count + 1
                elif subnet[subnet_id]["ip_version"] == 6:
                    v6_count = v6_count + 1
                else:
                    print >> sys.stderr, "Invalid ip_version"
                    sys.exit(1)
        if v4_count > 1:
            print "Network " + network_id + " has more than one ipv4 subnets"
            ret = 1
        if v6_count > 1:
            print "Network " + network_id + " has more than one ipv6 subnets"
            ret = 1
    return ret


def scan(ctx):
    load_partition(ctx.bigips)
    load_vip(ctx.bigips)
    load_lbaas_db(ctx.db)
    sys.exit(scan_network())


def step1(ctx):
    load_partition(ctx.bigips)
    load_vip(ctx.bigips)
    load_vs(ctx.bigips)
    load_lbaas_db(ctx.db)
    create_snat_selfip(ctx.bigips, ctx.neutron)
    create_snat_port(ctx.neutron)
    create_dedicated_snat_pool(ctx.bigips)
    tmsh_save(ctx.bigips)


def step1r(ctx):
    load_partition(ctx.bigips)
    load_vip(ctx.bigips)
    delete_snat_selfip(ctx.bigips)
    delete_dedicated_snat_pool(ctx.bigips)
    delete_snat_port(ctx.neutron)
    tmsh_save(ctx.bigips)


def step2(ctx):
    load_partition(ctx.bigips)
    load_vs(ctx.bigips)
    switch_to_dedicated_snat_pool(ctx.bigips)
    tmsh_save(ctx.bigips)


def step2r(ctx):
    load_partition(ctx.bigips)
    load_vs(ctx.bigips)
    switch_to_shared_snat_pool(ctx.bigips)
    tmsh_save(ctx.bigips)


def step3(ctx):
    load_partition(ctx.bigips)
    delete_shared_snat_pool(ctx.bigips)
    delete_legacy_snat_port(ctx.neutron)
    tmsh_save(ctx.bigips)


def main():

    parser = argparse.ArgumentParser(
        description="BIG-IP SNAT Migration Tool")
    parser.add_argument("--bigip", type=str, required=True,
                        help="BIG-IP hostname")
    parser.add_argument("--username", type=str, default="admin",
                        help="BIG-IP username")
    parser.add_argument("--password", type=str, required=True,
                        help="BIG-IP password")
    parser.add_argument("--db-connection", type=str, required=True,
                        help="Neutron DB connection string")
    parser.add_argument("migrate", nargs="?", type=str, default="no",
                        help="SNAT migration step")

    args = parser.parse_args()
    username = args.username
    password = args.password
    db_conn = args.db_connection
    conn = connect_neutron_db(db_conn)
    migrate = args.migrate

    hosts = []
    start = 0
    while start >= 0:
        split = args.bigip.find(",", start)
        if split > start:
            hosts.append(args.bigip[start:start+split])
            start = split + 1
        elif split < 0:
            hosts.append(args.bigip[start:])
            start = -1
        else:
            start = -1

    bigips = []
    for host in hosts:
        bigips.append(ManagementRoot(host, username, password))

    get_device_name(bigips)

    ctx = Context()
    ctx.db = conn
    ctx.bigips = bigips
    ctx.neutron = get_neutron_client()

    if migrate == "scan":
        scan(ctx)

    if migrate == "step1":
        step1(ctx)

    if migrate == "step1r":
        step1r(ctx)

    if migrate == "step2":
        step2(ctx)

    if migrate == "step2r":
        step2r(ctx)

    if migrate == "step3":
        step3(ctx)


if __name__ == "__main__":
    main()
